import os
import pandas as pd
import numpy as np
import joblib
from datetime import datetime, timedelta
import sys
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.svm import SVR
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt

# 添加项目根目录到系统路径
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))

from config.config import DATA_CONFIG, MODEL_CONFIG

class WeatherPredictor:
    """
    天气预测模型类，负责模型训练、评估和预测
    """
    def __init__(self):
        self.processed_data_path = os.path.join(os.path.dirname(__file__), '../../', DATA_CONFIG['processed_data_path'])
        self.model_save_path = os.path.join(os.path.dirname(__file__), '../../', MODEL_CONFIG['save_path'])
        self.test_size = MODEL_CONFIG['train']['test_size']
        self.random_state = MODEL_CONFIG['train']['random_state']
        self.cv = MODEL_CONFIG['train']['cv']
        self.forecast_days = MODEL_CONFIG['predict']['forecast_days']
        self.confidence_interval = MODEL_CONFIG['predict']['confidence_interval']
        
        # 确保模型保存目录存在
        os.makedirs(self.model_save_path, exist_ok=True)
        
        # 初始化模型字典
        self.models = {
            'linear': LinearRegression(),
            'ridge': Ridge(),
            'lasso': Lasso(),
            'svr': SVR(),
            'random_forest': RandomForestRegressor(random_state=self.random_state),
            'gradient_boosting': GradientBoostingRegressor(random_state=self.random_state)
        }
        
        # 初始化特征缩放器
        self.scaler = StandardScaler()
        
        # 初始化最佳模型
        self.best_model = None
        self.best_model_name = None
        self.feature_importance = None
    
    def load_data(self, data_type='historical'):
        """
        加载处理后的数据用于模型训练
        """
        file_map = {
            'historical': 'historical_weather.csv',
            'all': 'all_weather_data.csv'
        }
        
        if data_type not in file_map:
            print(f"未知的数据类型: {data_type}")
            return None
        
        file_path = os.path.join(self.processed_data_path, file_map[data_type])
        if not os.path.exists(file_path):
            print(f"文件不存在: {file_path}")
            return None
        
        try:
            df = pd.read_csv(file_path)
            
            # 转换时间列
            time_cols = ['timestamp', 'forecast_time']
            for col in time_cols:
                if col in df.columns:
                    df[col] = pd.to_datetime(df[col])
            
            return df
        except Exception as e:
            print(f"加载数据失败: {e}")
            return None
    
    def prepare_features(self, df, target_col='temp', city=None):
        """
        准备模型特征和目标变量
        """
        if df is None or df.empty:
            return None, None, None
        
        # 如果指定了城市，则筛选数据
        if city and 'city' in df.columns:
            df = df[df['city'] == city].copy()
            if df.empty:
                print(f"没有找到城市 {city} 的数据")
                return None, None, None
        
        # 确保目标列存在
        if target_col not in df.columns:
            print(f"目标列 {target_col} 不存在")
            return None, None, None
        
        # 选择特征列
        feature_cols = []
        
        # 时间特征
        if 'year' in df.columns: feature_cols.append('year')
        if 'month' in df.columns: feature_cols.append('month')
        if 'day' in df.columns: feature_cols.append('day')
        if 'hour' in df.columns: feature_cols.append('hour')
        if 'dayofweek' in df.columns: feature_cols.append('dayofweek')
        
        # 气象特征
        weather_features = ['humidity', 'pressure', 'wind_speed', 'clouds']
        for col in weather_features:
            if col in df.columns and col != target_col:
                feature_cols.append(col)
        
        # 确保有足够的特征
        if len(feature_cols) < 3:
            print("特征数量不足，无法训练模型")
            return None, None, None
        
        # 去除缺失值
        df_clean = df.dropna(subset=feature_cols + [target_col])
        if df_clean.empty:
            print("清洗后数据为空")
            return None, None, None
        
        # 准备特征和目标变量
        X = df_clean[feature_cols]
        y = df_clean[target_col]
        
        return X, y, feature_cols
    
    def train_models(self, X, y, feature_cols, city=None, target_col='temp'):
        """
        训练多个模型并选择最佳模型
        """
        if X is None or y is None:
            return None
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=self.test_size, random_state=self.random_state
        )
        
        # 特征缩放
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        # 训练并评估每个模型
        results = {}
        for name, model in self.models.items():
            try:
                # 训练模型
                model.fit(X_train_scaled, y_train)
                
                # 预测
                y_pred = model.predict(X_test_scaled)
                
                # 评估
                mse = mean_squared_error(y_test, y_pred)
                rmse = np.sqrt(mse)
                mae = mean_absolute_error(y_test, y_pred)
                r2 = r2_score(y_test, y_pred)
                
                results[name] = {
                    'model': model,
                    'mse': mse,
                    'rmse': rmse,
                    'mae': mae,
                    'r2': r2
                }
                
                print(f"模型 {name} 评估结果: RMSE={rmse:.4f}, MAE={mae:.4f}, R²={r2:.4f}")
            
            except Exception as e:
                print(f"训练模型 {name} 时出错: {e}")
        
        if not results:
            print("没有成功训练的模型")
            return None
        
        # 选择最佳模型（基于RMSE）
        self.best_model_name = min(results, key=lambda x: results[x]['rmse'])
        self.best_model = results[self.best_model_name]['model']
        
        print(f"最佳模型: {self.best_model_name}, RMSE={results[self.best_model_name]['rmse']:.4f}")
        
        # 获取特征重要性（如果模型支持）
        if hasattr(self.best_model, 'feature_importances_'):
            self.feature_importance = dict(zip(feature_cols, self.best_model.feature_importances_))
            print("特征重要性:")
            for feature, importance in sorted(self.feature_importance.items(), key=lambda x: x[1], reverse=True):
                print(f"{feature}: {importance:.4f}")
        
        # 保存最佳模型
        model_filename = f"{self.best_model_name}_{target_col}"
        if city:
            model_filename += f"_{city}"
        model_filename += ".joblib"
        
        model_path = os.path.join(self.model_save_path, model_filename)
        joblib.dump({
            'model': self.best_model,
            'scaler': self.scaler,
            'feature_cols': feature_cols,
            'target_col': target_col,
            'metrics': results[self.best_model_name],
            'feature_importance': self.feature_importance
        }, model_path)
        
        print(f"模型已保存至: {model_path}")
        
        return results
    
    def load_model(self, model_name, target_col='temp', city=None):
        """
        加载已训练的模型
        """
        model_filename = f"{model_name}_{target_col}"
        if city:
            model_filename += f"_{city}"
        model_filename += ".joblib"
        
        model_path = os.path.join(self.model_save_path, model_filename)
        if not os.path.exists(model_path):
            print(f"模型文件不存在: {model_path}")
            return False
        
        try:
            model_data = joblib.load(model_path)
            self.best_model = model_data['model']
            self.scaler = model_data['scaler']
            self.feature_cols = model_data['feature_cols']
            self.target_col = model_data['target_col']
            self.feature_importance = model_data.get('feature_importance')
            self.best_model_name = model_name
            
            print(f"成功加载模型: {model_path}")
            return True
        except Exception as e:
            print(f"加载模型失败: {e}")
            return False
    
    def predict(self, input_data, days=None):
        """
        使用训练好的模型进行预测
        input_data: 输入数据，可以是DataFrame或特征字典
        days: 预测未来天数，默认使用配置中的值
        """
        if self.best_model is None:
            print("没有训练好的模型，请先训练或加载模型")
            return None
        
        if days is None:
            days = self.forecast_days
        
        # 如果输入是DataFrame，提取特征
        if isinstance(input_data, pd.DataFrame):
            if all(col in input_data.columns for col in self.feature_cols):
                # 使用当前数据进行预测
                X = input_data[self.feature_cols].iloc[-1:]
            else:
                print(f"输入数据缺少必要的特征列: {self.feature_cols}")
                return None
        elif isinstance(input_data, dict):
            # 检查字典中是否包含所有特征
            if all(col in input_data for col in self.feature_cols):
                X = pd.DataFrame([input_data])
            else:
                print(f"输入数据缺少必要的特征列: {self.feature_cols}")
                return None
        else:
            print("不支持的输入数据类型")
            return None
        
        # 初始化结果列表
        predictions = []
        dates = []
        
        # 获取起始日期
        if isinstance(input_data, pd.DataFrame) and 'timestamp' in input_data.columns:
            start_date = input_data['timestamp'].iloc[-1]
        else:
            start_date = datetime.now()
        
        # 逐天预测
        current_date = start_date
        current_features = X.copy()
        
        for i in range(days):
            # 更新日期
            current_date = current_date + timedelta(days=1)
            dates.append(current_date)
            
            # 更新时间特征
            if 'year' in self.feature_cols:
                current_features['year'] = current_date.year
            if 'month' in self.feature_cols:
                current_features['month'] = current_date.month
            if 'day' in self.feature_cols:
                current_features['day'] = current_date.day
            if 'dayofweek' in self.feature_cols:
                current_features['dayofweek'] = current_date.weekday()
            
            # 特征缩放
            X_scaled = self.scaler.transform(current_features)
            
            # 预测
            pred = self.best_model.predict(X_scaled)[0]
            predictions.append(pred)
            
            # 更新特征（如果预测的是温度，可能会影响其他特征）
            if self.target_col in self.feature_cols:
                current_features[self.target_col] = pred
        
        # 创建结果DataFrame
        result_df = pd.DataFrame({
            'date': dates,
            self.target_col: predictions
        })
        
        return result_df
    
    def evaluate_model(self, X_test, y_test):
        """
        评估模型性能
        """
        if self.best_model is None:
            print("没有训练好的模型，请先训练或加载模型")
            return None
        
        # 特征缩放
        X_test_scaled = self.scaler.transform(X_test)
        
        # 预测
        y_pred = self.best_model.predict(X_test_scaled)
        
        # 计算评估指标
        mse = mean_squared_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        mae = mean_absolute_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        
        # 打印评估结果
        print(f"模型评估结果:")
        print(f"RMSE: {rmse:.4f}")
        print(f"MAE: {mae:.4f}")
        print(f"R²: {r2:.4f}")
        
        # 绘制预测值与实际值对比图
        plt.figure(figsize=(12, 6))
        plt.plot(y_test.values, label='实际值')
        plt.plot(y_pred, label='预测值')
        plt.title(f'{self.best_model_name} 模型预测结果')
        plt.xlabel('样本')
        plt.ylabel(self.target_col)
        plt.legend()
        plt.tight_layout()
        
        # 保存图表
        fig_path = os.path.join(self.model_save_path, f"{self.best_model_name}_{self.target_col}_evaluation.png")
        plt.savefig(fig_path)
        plt.close()
        
        print(f"评估图表已保存至: {fig_path}")
        
        return {
            'mse': mse,
            'rmse': rmse,
            'mae': mae,
            'r2': r2,
            'y_pred': y_pred
        }
    
    def train_for_all_cities(self, data_type='historical', target_col='temp'):
        """
        为所有城市训练模型
        """
        df = self.load_data(data_type)
        if df is None:
            return None
        
        if 'city' not in df.columns:
            print("数据中没有城市信息")
            return None
        
        cities = df['city'].unique()
        results = {}
        
        for city in cities:
            print(f"\n为城市 {city} 训练模型...")
            X, y, feature_cols = self.prepare_features(df, target_col=target_col, city=city)
            if X is not None and y is not None:
                model_results = self.train_models(X, y, feature_cols, city=city, target_col=target_col)
                if model_results:
                    results[city] = {
                        'best_model': self.best_model_name,
                        'rmse': model_results[self.best_model_name]['rmse']
                    }
        
        return results